import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
from fastai.vision.all import L, Callback
import os.path as osp


def plot_losses(learner, filename, skip_start=5, with_valid=True):
    plt.figure()
    plt.plot(list(range(skip_start, len(learner.recorder.losses))), learner.recorder.losses[skip_start:], label='train')
    if with_valid:
        idx = (np.array(learner.recorder.iters) < skip_start).sum()
        plt.plot(learner.recorder.iters[idx:], L(learner.recorder.values[idx:]).itemgot(1), label='valid')
        plt.legend()
    plt.savefig(filename)
    plt.close()


class PlotCallback(Callback):

    def __init__(self, model_path, model_id, fold_idx):
        self.model_path = model_path
        self.model_id = model_id
        self.fold_idx = fold_idx
        self.total_steps = 0

    def after_batch(self):
        if self.learn.training:
            self.total_steps = len(self.learn.dls.train)
            if self.learn.iter % (self.total_steps // 10) == 0:  # Only plot 10 times every epoch
                plot_losses(self.learn, osp.join(self.model_path, f'model_{self.model_id}_{self.fold_idx}_loss.jpg'))
